#include <iostream>
#include <string>
#include <vector>
#include <thread>
#include <functional> // For std::function
#include <memory>     // For std::unique_ptr

// Tracy Profiler header
#include <Tracy.hpp>

// CLI11 for command-line argument parsing
#include <CLI/CLI.hpp>

// Include your signature manager headers
#include "signature_manager/RSA.h"
#include "signature_manager/DSA.h"
#include "signature_manager/Rabin.h"
#include "signature_manager/ECDSA.h"
#include "signature_manager/EdDSA.h"
#include "signature_manager/Hash.h" // Assuming SHA256 is here
#include "signature_manager/signature_manager.h"
#include "signature_manager/data_conversion.h" // For base64_encode/decode, if needed
#include <TracyC.h>

// Forward declarations or proper includes for your specific types if not already in .h files
// For example, if secp256k1_tag, prime256v1_tag etc. are defined in separate files or within the same signature_manager.h
// Make sure all necessary headers for RSA_type_traits, DSA_type_traits, Rabin_type_traits, ECDSA_type_traits, EdDSA_type_traits are included.

// worker_function: Each thread will execute this function to perform signature/verification operations.
// It takes a factory function to create a SignatureManager instance, number of iterations,
// and the message to be signed.
void worker_function(
    std::function<std::unique_ptr<SignatureManager_base>()> manager_factory,
    int iterations,
    const PlainMessage& message_to_sign,
    int thread_id) {

    // Set a name for the Tracy thread, visible in the profiler UI
    std::string thread_name = "Worker Thread " + std::to_string(thread_id);
    TracyCSetThreadName(thread_name.c_str());

    // ZoneScoped captures the entire duration of this worker function
    ZoneScopedN("Worker_Function_Total");

    // Create a SignatureManager instance for this thread
    // Using unique_ptr for automatic memory management
    std::unique_ptr<SignatureManager_base> signature_manager = manager_factory();

    if (!signature_manager) {
        // This should ideally not happen if manager_factory is set up correctly
        std::cerr << "错误：无法在线程 " << thread_id << " 中创建签名管理器实例。\n";
        TracyMessageL("Failed to create SignatureManager in worker.");
        return;
    }

    // Generate key pair once per thread (outside the iteration loop)
    // This allows profiling the key generation overhead separately from sign/verify iterations.
    std::string public_key_str;
    std::string public_params_str;
    {
        ZoneScopedN("Key_Generation"); // Tracy zone for key generation
        signature_manager->generateKeyPair();
        public_key_str = signature_manager->get_publickey(); // Get public key
        public_params_str = signature_manager->get_publicparameters(); // Get public parameters (if any)
    }

    // Perform signature and verification iterations
    for (int i = 0; i < iterations; ++i) {
        // Tracy zone for signing operation
        std::string generated_signature_str;
        {
            ZoneScopedN("Sign_Operation");
            signature_manager->signMessage(message_to_sign, generated_signature_str);
        }

        // Tracy zone for verification operation
        bool result;
        {
            ZoneScopedN("Verify_Operation");
            // Set public key and parameters for verification.
            // Note: For algorithms like RSA/Rabin, public parameters might be empty.
            // For DSA, ECDSA, EdDSA, public parameters (e.g., curve details) are crucial.
            // The signature_manager instance already has its own internal public key/params from generateKeyPair.
            // If you were verifying a signature from *another* key pair, you'd need to set them here.
            // For profiling, we're verifying with the key pair generated by this same manager.
            signature_manager->set_publickey(public_key_str);
            signature_manager->set_publicparameters(public_params_str);
            result = signature_manager->verifySignature(message_to_sign, generated_signature_str);
        }

        if (!result) {
            // Log a message to Tracy if verification fails
            TracyMessageL("Verification failed for an iteration!");
            std::cerr << "警告：线程 " << thread_id << ", 迭代 " << i << " 签名验证失败。\n";
        }
    }
}


int main(int argc, char* argv[]) {
    // Tracy profiler initialization (optional, can be controlled by TRACY_ENABLE)
    // TracyClient; // If not using macro-based initialization

    CLI::App app{ "Digital Signature Performance Test with Tracy" };

    std::string signature_algorithm;
    int key_length = 0; // Default to 0, will be required for some algorithms
    std::string curve_name;
    int num_threads = 1;
    int iterations_per_thread = 1;
    std::string input_file_path; // Optional: path to a file to use as source data

    // Command line options
    app.add_option("-a,--algorithm", signature_algorithm, "签名算法 (RSA, DSA, Rabin, ECDSA, EdDSA)")->required();
    app.add_option("-l,--length", key_length, "密钥长度 (例如: 1024, 2048, 3072)，适用于 RSA, DSA, Rabin")->default_val(1024);
    app.add_option("-c,--curve", curve_name, "曲线名称 (例如: secp256k1, prime256v1, Ed25519, Ed448)，适用于 ECDSA, EdDSA")->default_val("");
    app.add_option("-t,--threads", num_threads, "并行线程数")->default_val(1);
    app.add_option("-i,--iterations", iterations_per_thread, "每个线程的签名/验签迭代次数")->default_val(100);
    app.add_option("-f,--file", input_file_path, "用于签名的输入文件路径。")->required();


    try {
        app.parse(argc, argv);
    } catch (const CLI::ParseError& e) {
        return app.exit(e); // Handles errors and --help automatically
    }

    // Input validation for algorithm-specific parameters
    if (signature_algorithm == "RSA" || signature_algorithm == "DSA" || signature_algorithm == "Rabin") {
        if (key_length != 1024 && key_length != 2048 && key_length != 3072) {
            std::cerr << "错误：对于 " << signature_algorithm << " 算法，无效的密钥长度。支持 1024, 2048, 3072。\n";
            return 1;
        }
    } else if (signature_algorithm == "ECDSA") {
        if (curve_name.empty()) {
            curve_name = "secp256k1"; // Default curve for ECDSA if not specified
            std::cout << "未指定 ECDSA 曲线，使用默认值 secp256k1。\n";
        } else if (curve_name != "secp256k1" && curve_name != "prime256v1" && curve_name != "secp384r1") {
            std::cerr << "错误：对于 ECDSA，无效的曲线名称。支持 secp256k1, prime256v1, secp384r1。\n";
            return 1;
        }
    } else if (signature_algorithm == "EdDSA") {
        if (curve_name.empty()) {
            curve_name = "Ed25519"; // Default curve for EdDSA if not specified
            std::cout << "未指定 EdDSA 曲线，使用默认值 Ed25519。\n";
        } else if (curve_name != "Ed25519" && curve_name != "Ed448") {
            std::cerr << "错误：对于 EdDSA，无效的曲线名称。支持 Ed25519, Ed448。\n";
            return 1;
        }
    } else {
        std::cerr << "错误：不支持的签名算法：" << signature_algorithm << "。\n";
        return 1;
    }

    // Prepare the dummy message for signing/verification
    PlainMessage message_data(new SHA256());
    if (input_file_path.empty()) {
        std::cerr << "错误：未提供输入文件。\n";
        return 1; // Error: No input file provided
    }

        try {
            message_data.read_file_to_bytes(input_file_path.c_str());
            std::cout << "从文件加载数据用于签名: " << input_file_path << "\n";
        } catch (const std::exception& e) {
            std::cerr << "错误：无法读取输入文件 " << input_file_path << ": " << e.what() << std::endl;
            return 1;
        }
    
    

    // Factory function to create SignatureManager instances based on selected algorithm and parameters
    std::function<std::unique_ptr<SignatureManager_base>()> manager_factory;

    if (signature_algorithm == "RSA") {
        if (key_length == 1024) {
            manager_factory = []() { return std::make_unique<SignatureManager<RSA_type_traits<1024>>>(); };
        } else if (key_length == 2048) {
            manager_factory = []() { return std::make_unique<SignatureManager<RSA_type_traits<2048>>>(); };
        } else if (key_length == 3072) {
            manager_factory = []() { return std::make_unique<SignatureManager<RSA_type_traits<3072>>>(); };
        }
    } else if (signature_algorithm == "DSA") {
        if (key_length == 1024) {
            manager_factory = []() { return std::make_unique<SignatureManager<DSA_type_traits<1024>>>(); };
        } else if (key_length == 2048) {
            manager_factory = []() { return std::make_unique<SignatureManager<DSA_type_traits<2048>>>(); };
        } else if (key_length == 3072) {
            manager_factory = []() { return std::make_unique<SignatureManager<DSA_type_traits<3072>>>(); };
        }
    } else if (signature_algorithm == "Rabin") {
        if (key_length == 1024) {
            manager_factory = []() { return std::make_unique<SignatureManager<Rabin_type_traits<1024>>>(); };
        } else if (key_length == 2048) {
            manager_factory = []() { return std::make_unique<SignatureManager<Rabin_type_traits<2048>>>(); };
        } else if (key_length == 3072) {
            manager_factory = []() { return std::make_unique<SignatureManager<Rabin_type_traits<3072>>>(); };
        }
    } else if (signature_algorithm == "ECDSA") {
        if (curve_name == "secp256k1") {
            manager_factory = []() { return std::make_unique<SignatureManager<ECDSA_type_traits<secp256k1_tag>>>(); };
        } else if (curve_name == "prime256v1") {
            manager_factory = []() { return std::make_unique<SignatureManager<ECDSA_type_traits<prime256v1_tag>>>(); };
        } else if (curve_name == "secp384r1") {
            manager_factory = []() { return std::make_unique<SignatureManager<ECDSA_type_traits<secp384r1_tag>>>(); };
        }
    } else if (signature_algorithm == "EdDSA") {
        if (curve_name == "Ed25519") {
            manager_factory = []() { return std::make_unique<SignatureManager<EdDSA_type_traits<Ed25519_tag>>>(); };
        } else if (curve_name == "Ed448") {
            manager_factory = []() { return std::make_unique<SignatureManager<EdDSA_type_traits<Ed448_tag>>>(); };
        }
    }

    if (!manager_factory) {
        std::cerr << "内部错误：无法为指定算法创建签名管理器工厂。\n";
        return 1;
    }

    std::cout << "开始性能测试...\n";
    std::cout << "算法: " << signature_algorithm << ", 密钥长度/曲线: " << (key_length ? std::to_string(key_length) : curve_name)
              << ", 线程数: " << num_threads << ", 每个线程迭代次数: " << iterations_per_thread << std::endl;

    // Create and launch threads
    std::vector<std::thread> threads;
    for (int i = 0; i < num_threads; ++i) {
        threads.emplace_back(worker_function, manager_factory, iterations_per_thread, std::cref(message_data), i);
    }

    // Wait for all threads to complete
    for (std::thread& t : threads) {
        if (t.joinable()) {
            t.join();
        }
    }

    std::cout << "所有线程完成，性能测试结束。\n";

    // TracyClient::ShutDown(); // If using explicit shutdown, typically handled by Tracy's global shutdown

    return 0;
}
